package cn.genmer.test.security.machinelearning.deeplearning4j.mnist.V1;

import org.deeplearning4j.datasets.iterator.impl.MnistDataSetIterator;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.DataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.factory.Nd4j;
 
import java.io.File;
import java.io.IOException;

public class MnistPredictor {
    private MultiLayerNetwork model;
 
    public MnistPredictor() {
        // 加载预训练好的LeNet模型
        File locationToSave = new File(MnistTrain.BASE_PATH +"/cnn_mnist_model.zip");
        System.out.println("预训练模型的路径：" + locationToSave.getAbsolutePath());
        try {
            model = ModelSerializer.restoreMultiLayerNetwork(locationToSave);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
 
    public int predict(int[][] inputImage) {
        // 获取输入参数的形状
        int[] shape = getShape(inputImage);
        // 创建INDArray对象
        INDArray array = Nd4j.create(shape);

        // 将多维数组的值复制到INDArray对象中
        for (int i = 0; i < shape[0]; i++) {
            for (int j = 0; j < shape[1]; j++) {
                array.putScalar(new int[]{i, j}, inputImage[i][j]);
            }
        }

        // 调整形状并进行数值计算
        INDArray input = array.reshape(1, 1, 28, 28).divi(255.0);
        // 对图像进行推理，并返回预测结果的索引
        INDArray output = model.output(input);
        INDArray indArray = Nd4j.argMax(output, 1);
        return indArray.getInt(0);
    }
    public void predict2() throws IOException {
        // 创建测试迭代器
        DataSetIterator testIter = new MnistDataSetIterator(32, false, 12345);

// 获取一个批次的测试样本
        DataSet testData = testIter.next();

// 将测试样本的特征输入到模型中进行预测
        INDArray features = testData.getFeatures();
        INDArray predictions = model.output(features);

// 解析预测结果和标签
        int[] predictedLabels = predictions.argMax(1).toIntVector();
        int[] trueLabels = testData.getLabels().argMax(1).toIntVector();

// 输出预测结果和标签
        for (int i = 0; i < predictedLabels.length; i++) {
            System.out.println("预测结果：" + predictedLabels[i] + "，真实标签：" + trueLabels[i]);
        }
    }

    public static int[] getShape(int[][] array) {
        int firstDimensionSize = array.length;
        int secondDimensionSize = array[0].length;

        return new int[]{firstDimensionSize, secondDimensionSize};
    }
    public static int getRank(int[][] array) {
        // 获取第一个维度大小
        int firstDimensionSize = array.length;

        // 获取第二个维度大小
        int secondDimensionSize = array[0].length;

        // 返回维度数量
        if (firstDimensionSize > 0 && secondDimensionSize > 0) {
            return 2;
        } else if (firstDimensionSize > 0 || secondDimensionSize > 0) {
            return 1;
        } else {
            return 0;
        }
    }

    private INDArray prepareInputData(int[][] input) {
        int height = input.length;
        int width = input[0].length;
        float[][] floatInput = new float[height][width];

        for (int i = 0; i < height; i++) {
            for (int j = 0; j < width; j++) {
                floatInput[i][j] = input[i][j] / 255f;
            }
        }

        INDArray inputArray = Nd4j.create(floatInput).reshape(1, height * width);
        return inputArray;
    }
}